mod api_operation;
mod sse_invoke_method;
mod custom_jwt;
use crate::sse_invoke_method::sse_invoke::constant_value::{API_KEY};

use std::env;
use std::fs;
use std::io::Read;

#[allow(unused_imports)]
use std::io::{self, Write};
#[allow(unused_imports)]
use std::path::Path;
use std::collections::HashMap;
use regex::Regex;

fn remove_markdown_code_block(text: &str) -> String {
    // 创建正则表达式来匹配代码块标记
    let re = Regex::new(r"```.*?\n").unwrap();
    let text = re.replace_all(text, "");

    // 去除剩余的代码块结束标记
    let re = Regex::new(r"```").unwrap();
    let text = re.replace_all(&text, "");

    text.to_string()
}

fn traverse_dir(dir_path:&str, header_files: &mut HashMap<String, String>, source_files: &mut HashMap<String, String>)->std::io::Result<()>{
    let dir =fs::read_dir(dir_path)?;
    for entry in dir{
        let entry=entry?;
        let path=entry.path();
        if path.is_file(){
            //println!("File name:{:?}",path.file_name().unwrap());
            if let Some(extension) = path.extension() {
                if extension == "h" {
                    // 读取头文件内容
                    if let Ok(content) = fs::read_to_string(&path) {
                        let filename = path.file_stem().unwrap().to_str().unwrap().to_string();
                        header_files.insert(filename.clone(), content);
                    }
                } else if extension == "c" {
                    // 读取C源文件内容
                    if let Ok(content) = fs::read_to_string(&path) {
                        let filename = path.file_stem().unwrap().to_str().unwrap().to_string();
                        source_files.insert(filename.clone(), content);
                    }
                }
            }
        }
        else if path.is_dir(){
            let _= traverse_dir(&path.to_string_lossy(), header_files, source_files).unwrap();
        }
    }
    Ok(())
}

fn translate_prompt() -> Vec<String> {
    let mut ret = Vec::new();
    // 指定目录
    let directory: &str = "../Input/01-Primary/src";

    // 读取目录中的所有文件
    let mut header_files = HashMap::new();
    let mut source_files = HashMap::new();

    let _= traverse_dir(directory, &mut header_files, &mut source_files).unwrap();

    // 处理文件对
    for (filename, header_content) in &header_files {
        if let Some(source_content) = source_files.get(filename) {
            // 生成转译代码的prompt
            let prompt = format!("Translate the following C code to Rust: {}, ", source_content);
            ret.push(prompt);
        } 
    }

    ret
}


#[derive(Debug)]
pub struct C2RustGLM {
    chatglm_response: String,
}

impl C2RustGLM {
    pub async fn new() -> Self {
        C2RustGLM {
            chatglm_response: String::new(),
        }
    }

    pub async fn sse_invoke_calling(&mut self, api_key: &str, user_input: &str) {
        let mut require_calling = "SSE".to_string();
        let mut input = String::new();

        
        let api_key_instance = api_operation::APIKeys::get_instance(api_key);
        let jwt_creator =
            custom_jwt::CustomJwt::new(api_key_instance.get_user_id(), api_key_instance.get_user_secret());
        let jwt = jwt_creator.create_jwt();

        let jwt_to_verify = jwt.clone();
        let is_valid = jwt_creator.verify_jwt(&jwt_to_verify);
    
        let sse_call = sse_invoke_method::ReceiveSSEInvokeModelOnlyText::new(&jwt, &user_input.trim());

        self.chatglm_response = match sse_call.await.get_response_message() {
            Some(message) => message.to_string(), // Return the message as String
            None => "Error: Unable to get SSE response.".to_string(),
        }
    }

    pub fn get_ai_response(&self) -> String {
        self.chatglm_response.clone()
    }
}

#[tokio::main]
async fn main()
{
    let mut c_2_rust_glm = C2RustGLM::new().await;
    let prompts = translate_prompt();
    println!("prompts: {:?}", prompts.len());
    for prompt in prompts {
        c_2_rust_glm.sse_invoke_calling(API_KEY, prompt.as_str()).await;
        println!("{}", c_2_rust_glm.get_ai_response());
    }
    
}